package cn.iocoder.yudao.framework.datapermission.v2;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ObjUtil;
import cn.hutool.core.util.StrUtil;
import cn.iocoder.yudao.framework.common.exception.KnownServiceException;
import cn.iocoder.yudao.framework.common.util.reflect.ReflectUtils;
import cn.iocoder.yudao.framework.common.util.spring.SpringUtils;
import com.alibaba.ttl.TransmittableThreadLocal;
import com.baomidou.mybatisplus.annotation.TableField;
import com.baomidou.mybatisplus.annotation.TableId;
import com.baomidou.mybatisplus.core.metadata.TableInfo;
import com.baomidou.mybatisplus.core.metadata.TableInfoHelper;
import com.baomidou.mybatisplus.extension.plugins.handler.MultiDataPermissionHandler;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.Alias;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.LongValue;
import net.sf.jsqlparser.expression.StringValue;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.IsNullExpression;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;

import java.io.Serializable;
import java.lang.reflect.Field;
import java.util.*;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;

import static cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils.skipTenantPermissionCheck;

/**
 * <pre>
 * OOoO0OOoO0OOOooo0oOOOO0OOOOO0oooOO0ooOOO0Ooooo0OOOOo0ooooO0OOooo0Ooooo0OOOOO
 *  数据权限管理器，参阅 <a href="https://gitee.com/dibo_software/diboot/blob/develop-v3/diboot-core/src/main/java/com/diboot/core/handler/DataAccessControlHandler.java">diboot#DataAccessControlHandler</a> 类及 <a href="https://gitee.com/dibo_software/diboot/blob/develop-v3/diboot-core/src/main/java/com/diboot/core/data/access/DataAccessAnnoCache.java">diboot#DataAccessAnnoCache</a> 类
 * OOoO0OOoO0OOOooo0oOOOO0OOOOO0oooOO0ooOOO0Ooooo0OOOOo0ooooO0OOooo0Ooooo0OOOOO
 * </pre>
 *
 * @author mazc@dibo.ltd
 * @author 山野羡民（1032694760@qq.com）
 * @since 2025/04/06
 */
@SuppressWarnings("unused")
@Slf4j
public class DataPermissionManager implements MultiDataPermissionHandler {
    private static final Map<String, Map<String, String>> ANNOTATION_CACHE = new ConcurrentHashMap<>();
    private static final Map<Class<?>, DataPermissionScopeProvider> PROVIDER_CACHE = new LinkedHashMap<>();
    private static final ThreadLocal<Boolean> IGNORE_DATA_PERMISSION = new TransmittableThreadLocal<>();

    @Override
    public Expression getSqlSegment(Table table, Expression where, String mappedStatementId) {
        // 特殊：跨租户访问
        if (skipTenantPermissionCheck()) {
            return null;
        }
        // 如果忽略此来源
        Boolean ignoreQuery = IGNORE_DATA_PERMISSION.get();
        if (ignoreQuery != null) {
            IGNORE_DATA_PERMISSION.remove();
            log.debug("获取到 IGNORE_DATA_PERMISSION 变量并移除：{}", ignoreQuery);
            if (ignoreQuery) {
                return null;
            }
        }
        TableInfo tableInfo = TableInfoHelper.getTableInfo(StrUtil.removeAny(table.getName(), "`", "\"", "["));
        // 无权限检查点注解，不处理
        if (tableInfo == null || tableInfo.getEntityType() == null || !hasDataPermissionCheckpoint(tableInfo.getEntityType())) {
            return null;
        }
        return buildSqlSegment(table, tableInfo.getEntityType());
    }

    private Expression buildSqlSegment(Table mainTable, Class<?> entityClass) {
        return getDataPermissionTableColumns(entityClass).entrySet().stream().map(entry -> {
            DataPermissionScopeProvider scopeProvider = getDataPermissionScopeProvider(entityClass);
            if (scopeProvider == null) {
                log.warn("未获取到({})类的数据范围提供者的实现，请检查({})实现类是否正确实例化并指定作用于此实体！", entityClass.getSimpleName(), DataPermissionScopeProvider.class.getName());
                throw new KnownServiceException("无法从上下文中获取数据范围提供者的实现");
            }
            List<? extends Serializable> scopeIds = scopeProvider.getScopeIds(entityClass, entry.getKey());
            if (scopeIds == null) {
                // 未提供数据范围编号，则忽略数据权限处理
                return null;
            }
            String entryValue = entry.getValue();
            Alias tableAlias = mainTable.getAlias();
            if (tableAlias != null) {
                entryValue = tableAlias.getName() + "." + entryValue;
            }
            if (scopeIds.isEmpty()) {
                // 提供的数据范围编号为空，说明没该数据权限，则返回空条件
                return new IsNullExpression().withLeftExpression(new Column(entryValue));
            }
            if (scopeIds.size() == 1) {
                // 单个数据范围编号，则直接返回条件
                EqualsTo equalsTo = new EqualsTo();
                equalsTo.setLeftExpression(new Column(entryValue));
                if (scopeIds.get(0) instanceof Long) {
                    equalsTo.setRightExpression(new LongValue((Long) scopeIds.get(0)));
                } else {
                    equalsTo.setRightExpression(new StringValue(ObjUtil.defaultIfNull(scopeIds.get(0), "").toString()));
                }
                return equalsTo;
            }
            String condition = entryValue + " IN ";
            if (scopeIds.get(0) instanceof Number) {
                condition += "(" + StrUtil.join(", ", scopeIds) + ")";
            } else {
                condition += "('" + StrUtil.join("', '", scopeIds) + "')";
            }
            try {
                return CCJSqlParserUtil.parseCondExpression(condition);
            } catch (JSQLParserException e) {
                log.warn("解析 SQL 片段的查询条件异常: {}", condition, e);
            }
            return null;
        }).filter(Objects::nonNull).reduce(AndExpression::new).orElse(null);
    }

    /**
     * 忽略数据权限，执行对应的逻辑
     *
     * @param runnable 逻辑
     */
    public static void executeIgnore(Runnable runnable) {
        IGNORE_DATA_PERMISSION.set(true);
        try {
            runnable.run();
        } finally {
            IGNORE_DATA_PERMISSION.remove();
        }
    }

    /**
     * 忽略数据权限，执行对应的逻辑
     *
     * @param callable 逻辑
     * @return 执行结果
     */
    @SneakyThrows
    public static <T> T executeIgnore(Callable<T> callable) {
        IGNORE_DATA_PERMISSION.set(true);
        try {
            return callable.call();
        } finally {
            IGNORE_DATA_PERMISSION.remove();
        }
    }

    public synchronized DataPermissionScopeProvider getDataPermissionScopeProvider(Class<?> entityClass) {
        initDataPermissionScopeProvider();
        DataPermissionScopeProvider provider = PROVIDER_CACHE.get(entityClass);
        if (provider != null) {
            log.warn("获取到 {} 类的数据范围控制实现: {}", entityClass.getName(), provider.getClass().getName());
        }
        return provider;
    }

    private void initDataPermissionScopeProvider() throws KnownServiceException {
        List<DataPermissionScopeProvider> providers = SpringUtils.getBeans(DataPermissionScopeProvider.class);
        if (CollUtil.isEmpty(providers)) {
            return;
        }
        for (DataPermissionScopeProvider provider : providers) {
            List<Class<?>> entityClasses = provider.getEntityClasses();
            if (CollUtil.isEmpty(entityClasses)) {
                continue;
            }
            for (Class<?> aClass : entityClasses) {
                if (PROVIDER_CACHE.containsKey(aClass)) {
                    throw new KnownServiceException("多个数据权限范围提供者实现类作用于数据实体(%s)，请检查！".formatted(aClass.getName()));
                }
                PROVIDER_CACHE.put(aClass, provider);
                log.info("缓存数据实体({})与数据权限范围提供者实现类({})的对应关系", aClass.getName(), provider.getClass().getName());
            }
        }
    }

    public static boolean hasDataPermissionCheckpoint(Class<?> entityClass) {
        return !ANNOTATION_CACHE.computeIfAbsent(entityClass.getName(), k -> initDataPermissionCheckpoint(entityClass)).isEmpty();
    }

    public static Map<String, String> getDataPermissionTableColumns(Class<?> entityClass) {
        return ANNOTATION_CACHE.computeIfAbsent(entityClass.getName(), k -> initDataPermissionCheckpoint(entityClass));
    }

    private static Map<String, String> initDataPermissionCheckpoint(Class<?> entityClass) {
        List<Field> fieldList = ReflectUtils.extractFields(entityClass, DataPermissionCheckpoint.class);
        if (CollUtil.isNotEmpty(fieldList)) {
            Map<String, String> results = new HashMap<>();
            for (Field field : fieldList) {
                results.put(field.getName(), getTableColumnName(field));
            }
            return results;
        }
        return Collections.emptyMap();
    }

    private static String getTableColumnName(Field field) {
        String columnName = null;
        if (field.isAnnotationPresent(TableField.class)) {
            columnName = field.getAnnotation(TableField.class).value();
        } else if (field.isAnnotationPresent(TableId.class)) {
            columnName = field.getAnnotation(TableId.class).value();
        }
        return Optional.ofNullable(columnName).orElse(StrUtil.toUnderlineCase(field.getName()));
    }

}
